"""
版本号：1.0
日期：2025/4/9
描述：MCP客户端 - 支持多服务器连接的智能助手
"""

import asyncio
import os
import json
import sys
from typing import Optional, Dict, List, Tuple, Any
from contextlib import AsyncExitStack
from openai import OpenAI
from mcp import ClientSession, StdioServerParameters, Tool
from mcp.client.stdio import stdio_client

from utils import initialize_mcp_json, initialize_config, get_server_info


class MCPClient:
    def __init__(self):
        """初始化 MCP 客户端"""
        self.config = initialize_config()
        self.exit_stack = AsyncExitStack()
        
        # 从环境变量读取配置
        self.openai_api_key = self.config.get('ai', 'api_key')  # 读取 OpenAI API Key
        self.base_url = self.config.get('ai', 'base_url')  # 读取 BASE URL
        self.model = self.config.get('ai', 'model')  # 读取模型名称
        self.prompt = self.config.get('ai', 'prompt')  # 读取提示词
        self.temperature = self.config.getfloat('ai', 'temperature')  # 读取温度
        self.max_tokens = self.config.getint('ai', 'max_tokens')  # 读取最大令牌数
        self.server_list = initialize_mcp_json()  # 初始化 MCP 服务器列表

        if not self.openai_api_key:
            raise ValueError("❌ 未找到 OpenAI API Key，请在 .env 文件中设置 OPENAI_API_KEY")

        # 初始化 OpenAI 客户端
        self.client = OpenAI(api_key=self.openai_api_key, base_url=self.base_url)
        
        # 多会话存储
        self.sessions: Dict[str, ClientSession] = {}  # {server_name: session_object}
        self.server_transports: Dict[str, Tuple] = {}  # {server_name: (stdio, write)}
        self.tool_to_server_map: Dict[str, str] = {}  # {tool_name: server_name}
        self.all_tools: List[Tool] = []  # 存储所有工具的列表
        
        # 处理重名工具的计数器
        self.tool_name_counter: Dict[str, int] = {}  # 记录工具名称出现次数

    async def connect_to_server(self):
        """连接到 MCP 服务器并列出可用工具"""
        all_tools = []
        
        if not self.server_list:
            print("⚠️ 没有找到任何服务器配置")
            return
            
        for server in self.server_list:
            try:
                # 根据脚本类型选择执行命令
                command, args, env, name = get_server_info(server)
                if not command:
                    print(f"⚠️ 服务器 {name} 的配置无效，跳过连接")
                    continue
                    
                print(f"🔄 正在连接到服务器 {name}...")
                
                server_params = StdioServerParameters(
                    command=command,
                    args=args,
                    env=env
                )

                # 启动 MCP 服务器并建立通信
                stdio_transport = await self.exit_stack.enter_async_context(
                    stdio_client(server_params)
                )
                stdio, write = stdio_transport
                self.server_transports[name] = (stdio, write)
                
                # 创建并存储会话
                session = await self.exit_stack.enter_async_context(
                    ClientSession(stdio, write)
                )
                await session.initialize()
                self.sessions[name] = session

                # 列出 MCP 服务器上的工具
                response = await session.list_tools()
                server_tools = response.tools
                
                # 处理工具，解决重名问题
                for tool in server_tools:
                    # 检查是否已存在同名工具
                    original_name = tool.name
                    if original_name in self.tool_to_server_map:
                        # 如果有重名，添加服务器名称作为前缀以区分
                        self.tool_name_counter[original_name] = self.tool_name_counter.get(original_name, 1) + 1
                        count = self.tool_name_counter[original_name]
                        
                        # 如果是第一个发现的重名工具，需要特殊处理
                        if count == 2:
                            # 找到第一个相同名称的工具，并修改其名称
                            for i, existing_tool in enumerate(all_tools):
                                if existing_tool.name == original_name:
                                    first_server = self.tool_to_server_map[original_name]
                                    # 修改工具名称
                                    new_name = f"{first_server}_{original_name}"
                                    # 更新映射
                                    self.tool_to_server_map.pop(original_name)
                                    self.tool_to_server_map[new_name] = first_server
                                    # 更新工具名称
                                    all_tools[i].name = new_name
                                    break
                        
                        # 修改当前工具的名称
                        new_tool_name = f"{name}_{original_name}"
                        tool.name = new_tool_name
                        self.tool_to_server_map[new_tool_name] = name
                    else:
                        # 如果没有重名，直接添加映射
                        self.tool_to_server_map[original_name] = name
                
                all_tools.extend(server_tools)
                print(f"✅ 已连接到服务器 {name}，支持以下工具:", [tool.name for tool in server_tools])
            except Exception as e:
                print(f"❌ 连接服务器 {name} 失败: {str(e)}")
                import traceback
                print(traceback.format_exc())
        
        if not all_tools:
            print("⚠️ 没有从任何服务器获取到工具")
            return
            
        self.all_tools = all_tools
        print(f"\n🔧 所有服务器支持的工具汇总 ({len(all_tools)} 个):", [tool.name for tool in all_tools])

    async def process_query(self, query: str) -> str:
        """
        使用大模型处理查询并调用可用的 MCP 工具 (Function Calling)
        """
        messages = [{"role": "system",
                     "content": self.prompt},
                    {"role": "user", "content": query}]

        max_iterations = 10  # 防止无限循环
        iteration_count = 0

        while iteration_count < max_iterations:
            iteration_count += 1
            
            # 准备所有服务器的工具列表
            available_tools = [{
                "type": "function",
                "function": {
                    "name": tool.name,
                    "description": tool.description,
                    "input_schema": tool.inputSchema
                }
            } for tool in self.all_tools]

            # 调用 OpenAI API
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                tools=available_tools,
                temperature=self.temperature,
                max_tokens=self.max_tokens
            )

            # 处理返回的内容
            content = response.choices[0]
            
            # 如果不需要工具调用，直接返回内容
            if content.finish_reason != "tool_calls" or not content.message.tool_calls:
                return content.message.content
                
            # 添加模型消息到历史
            messages.append(content.message.model_dump())
            
            # 处理所有工具调用
            for tool_call in content.message.tool_calls:
                try:
                    tool_name = tool_call.function.name
                    tool_args = json.loads(tool_call.function.arguments)
                    
                    # 确定工具所在的服务器
                    server_name = self.tool_to_server_map.get(tool_name)
                    if not server_name:
                        error_msg = f"找不到工具 {tool_name} 所属的服务器"
                        print(f"\n⚠️ {error_msg}")
                        # 添加错误消息到历史
                        messages.append({
                            "role": "tool",
                            "content": error_msg,
                            "tool_call_id": tool_call.id,
                        })
                        continue
                        
                    # 获取对应服务器的会话
                    session = self.sessions.get(server_name)
                    if not session:
                        error_msg = f"服务器 {server_name} 的会话不存在"
                        print(f"\n⚠️ {error_msg}")
                        messages.append({
                            "role": "tool",
                            "content": error_msg,
                            "tool_call_id": tool_call.id,
                        })
                        continue
                    
                    print(f"\n\n[在服务器 {server_name} 上调用工具 {tool_name}，参数: {tool_args}]\n\n")
                    
                    # 执行工具
                    result = await session.call_tool(tool_name, tool_args)
                    
                    # 将结果存入消息历史
                    messages.append({
                        "role": "tool",
                        "content": result.content[0].text,
                        "tool_call_id": tool_call.id,
                    })
                except Exception as e:
                    error_msg = f"工具 {tool_name} 调用失败: {str(e)}"
                    print(f"\n⚠️ {error_msg}")
                    messages.append({
                        "role": "tool",
                        "content": error_msg,
                        "tool_call_id": tool_call.id,
                    })
            
        # 如果达到最大迭代次数，返回当前结果
        return "达到最大工具调用次数。请尝试简化您的请求或分多次提问。"

    async def chat_loop(self):
        """运行交互式聊天循环"""
        if not self.sessions:
            print("\n⚠️ 没有连接到任何服务器，无法开始对话")
            return
            
        server_names = list(self.sessions.keys())
        print(f"\n🤖 MCP 客户端已启动！已连接到以下服务器: {', '.join(server_names)}")
        print(f"可用工具总数: {len(self.all_tools)}")
        print("输入 'quit' 退出, 输入 'tools' 查看所有可用工具")
        
        while True:
            try:
                query = input("\n你: ").strip()
                if not query:
                    continue
                    
                if query.lower() == 'quit':
                    print("👋 再见！")
                    break
                    
                if query.lower() == 'tools':
                    print("\n🔧 可用工具列表:")
                    for tool in self.all_tools:
                        server_name = self.tool_to_server_map.get(tool.name, "未知")
                        print(f"{tool.name} (服务器: {server_name}): {tool.description}")
                    continue
                
                print("\n处理中...")
                response = await self.process_query(query)  # 发送用户输入到 OpenAI API
                print(f"\n🤖 OpenAI: {response}")
            except Exception as e:
                print(f"\n⚠️ 发生错误: {str(e)}")
                import traceback
                print(traceback.format_exc())

    async def cleanup(self):
        """清理资源"""
        try:
            print("\n🧹 清理资源...")
            # 关闭所有会话和连接
            await self.exit_stack.aclose()
            print("✅ 资源清理完成")
        except Exception as e:
            print(f"⚠️ 资源清理时发生错误: {str(e)}")
            import traceback
            print(traceback.format_exc())


async def main():
    """主函数"""
    print("🚀 启动 MCP 客户端...")
    
    try:
        client = MCPClient()
        print("🔌 正在连接到 MCP 服务器...")
        try:
            await client.connect_to_server()
            await client.chat_loop()
        except Exception as e:
            print(f"\n❌ 程序运行时发生错误: {str(e)}")
            import traceback
            print(traceback.format_exc())
    except Exception as e:
        print(f"\n❌ 初始化 MCP 客户端时发生错误: {str(e)}")
        import traceback
        print(traceback.format_exc())
    finally:
        # 确保客户端存在并且可以清理
        if 'client' in locals():
            await client.cleanup()
    
    print("🏁 程序已结束")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n\n👋 程序被用户中断")
    except Exception as e:
        print(f"\n❌ 程序崩溃: {str(e)}")
        import traceback
        print(traceback.format_exc())
